import collections
import concurrent.futures
import datetime
import functools
import json
import os
import logging
import time
import traceback

import bokeh.layouts
import bokeh.plotting
import pandas as pd
import pytz
from absl import app, flags
from bokeh.plotting import Figure
from bokeh.models import (
    BooleanFilter, Button, CheckboxGroup, ColumnDataSource, CDSView, CustomJS, 
    DataTable, DatePicker, DateFormatter, Div, Dropdown, GroupFilter, HoverTool,
    Legend, LegendItem, MultiChoice, RangeTool, Span, Select, TableColumn, TextInput)
from bokeh.server.server import Server
from tornado.ioloop import IOLoop

import coin.proto.coin_feed_pb2 as feedpb
import coin.proto.coin_request_pb2 as cr_pb2
import coin.strategy.mm.tool.archive_base as abase
from coin.base.param_util import to_list
from coin.base.query_util import query_pta
from coin.base.datetime_util import to_timestamp_int
from coin.exchange.kr_rest.product.product_impl import generate_product_from_str2
from coin.support.pta.app.order_event_dumper import get_order_event_data_dir


FLAGS = flags.FLAGS
_DataSourceInfo = collections.namedtuple('_DataSourceInfo', ['name', 'row_idx'])


def _str_to_datetime(dt_str):
  func = lambda x: datetime.datetime.strptime(dt_str, x).replace(tzinfo=pytz.UTC)
  try:
    return func('%Y-%m-%d %H:%M:%S')
  except:
    return func('%Y-%m-%d %H:%M:%S.%f')


def _parse_feed_type(feed_type):
  return None if feed_type == '' else feed_type


def _feed_display_func(iterator, feed_type):
  if feed_type == 'FEED_TYPE_BOOK':
    columns = ['timestamp', 'bid0p', 'bid0q', 'ask0p', 'ask0q', 'midp']
  elif feed_type == 'FEED_TYPE_TRADE':
    columns = ['timestamp', 'tradep', 'tradeq', 'trade_side']
  row_idx = iterator.row_idx
  row = pd.Series([None] * len(columns), index=columns)
  if row_idx is not None:
    # iloc doens't include end
    data = iterator.data.iloc[:row_idx+1]
    last_valid_idx = data[data['feed_type'] == feed_type].last_valid_index()
    if last_valid_idx is not None:
     row = data.loc[last_valid_idx, columns]
  text_list = []
  for column in columns:
    value = row[column]
    if column in ('event_timestamp', 'timestamp'):
      column = 'time'
    if pd.isnull(value):
      text = ''
    else:
      if column == 'time':
        text = datetime.datetime.utcfromtimestamp(value / 10**9).strftime(
            '%m-%d %H:%M:%S.%f')
      else:
        text = value
    text_list.append('%s: %s' % (column, text))
  return ', '.join(text_list)


def _order_display_func(iterator, columns, search_valid_column='type'):
  columns = to_list(columns)
  row_idx = iterator.row_idx
  row = pd.Series([None] * len(columns), index=columns)
  if row_idx is not None:
    # iloc doens't include end
    data = iterator.data.iloc[:row_idx+1]
    last_valid_idx = data[search_valid_column].last_valid_index()
    if last_valid_idx is not None:
     row = data.loc[last_valid_idx, columns]
  text_list = []
  for column in columns:
    value = row[column]
    if column in ('event_timestamp', 'timestamp'):
      column = 'time'
    if pd.isnull(value):
      text = ''
    else:
      if column == 'time':
        text = datetime.datetime.utcfromtimestamp(value / 10**9).strftime(
            '%m-%d %H:%M:%S.%f')
      else:
        text = value
    text_list.append('%s: %s' % (column, text))
  return ', '.join(text_list)


def _gen_order_sequence_diagram(
    *, source, selected_source, feed_sources, feed_type, plot_height, plot_width,
    start_time, end_time, tools):
  price_displayer = DataSourceDisplayer()
  columns = ['event_timestamp', 'type', 'order_side', 'order_price', 'order_qty',
             'fill_price', 'fill_qty']
  price_displayer.add_display_info(
      source,
      [functools.partial(_order_display_func, columns=columns)])
  price_fig = bokeh.plotting.figure(
      plot_height=plot_height, plot_width=plot_width, x_range=(start_time, end_time),
      x_axis_type='datetime', tools=tools, output_backend='webgl')
  muted_alpha = 0.2
  buy_filter = GroupFilter(column_name='order_sign', group='BUY_ORDER')
  sell_filter = GroupFilter(column_name='order_sign', group='SELL_ORDER')
  submit_filter = GroupFilter(column_name='type', group='ORDER_SUBMITTED')
  fill_filter = GroupFilter(column_name='type', group='ORDER_FILLED')
  cancel_filter = GroupFilter(column_name='type', group='CANCEL_CONFIRMED')
  buy_order_price_view = CDSView(source=source, filters=[buy_filter, submit_filter])
  sell_order_price_view = CDSView(source=source, filters=[sell_filter,  submit_filter])
  buy_fill_price_view = CDSView(source=source, filters=[buy_filter, fill_filter])
  sell_fill_price_view = CDSView(source=source, filters=[sell_filter, fill_filter])
  buy_cancel_view = CDSView(source=source, filters=[buy_filter, cancel_filter])
  sell_cancel_view = CDSView(source=source, filters=[sell_filter, cancel_filter])
  buy_order_price_circle = price_fig.circle(
      x='event_time', y='order_price', source=source, view=buy_order_price_view,
      line_color='green', fill_color='white', muted_alpha=muted_alpha)
  sell_order_price_circle = price_fig.circle(
      x='event_time', y='order_price', source=source, view=sell_order_price_view,
      line_color='red', fill_color='white', muted_alpha=muted_alpha)
  buy_fill_price_circle = price_fig.circle(
      x='event_time', y='fill_price', source=source, view=buy_fill_price_view,
      line_color='green', fill_color='green', muted_alpha=muted_alpha)
  sell_fill_price_circle = price_fig.circle(
      x='event_time', y='fill_price', source=source, view=sell_fill_price_view,
      line_color='red', fill_color='red', muted_alpha=muted_alpha)
  buy_fill_price_segment = price_fig.segment(
      x0='pending_time', y0='order_price', x1='event_time', y1='fill_price',
      source=source, view=buy_fill_price_view, line_color='green', line_dash='dotted',
      muted_alpha=muted_alpha)
  sell_fill_price_segment = price_fig.segment(
      x0='pending_time', y0='order_price', x1='event_time', y1='fill_price',
      source=source, view=sell_fill_price_view, line_color='red', line_dash='dotted',
      muted_alpha=muted_alpha)
  buy_cancel_x = price_fig.x(
      x='event_time', y='order_price', source=source, view=buy_cancel_view,
      line_color='green', muted_alpha=muted_alpha)
  sell_cancel_x = price_fig.x(
      x='event_time', y='order_price', source=source, view=sell_cancel_view,
      line_color='red', muted_alpha=muted_alpha)
  buy_cancel_segment = price_fig.segment(
      x0='pending_time', y0='order_price', x1='event_time', y1='order_price',
      source=source, view=buy_cancel_view, line_color='green', line_dash='dotted',
      muted_alpha=muted_alpha)
  sell_cancel_segment = price_fig.segment(
      x0='pending_time', y0='order_price', x1='event_time', y1='order_price',
      source=source, view=sell_cancel_view, line_color='red', line_dash='dotted',
      muted_alpha=muted_alpha)
  order_legend_items = [
      ('buy_order', [buy_order_price_circle]),
      ('selll_order' , [sell_order_price_circle]),
      ('buy_fill' , [buy_fill_price_circle, buy_fill_price_segment]),
      ('sell_fill' , [sell_fill_price_circle, sell_fill_price_segment]),
      ('buy_cancel' , [buy_cancel_x, buy_cancel_segment]),
      ('sell_cancel' , [sell_cancel_x, sell_cancel_segment]),
  ]
  feed_legend_items = []
  for product, feed_source in feed_sources.items():
    if feed_type == 'bbo':
      price_displayer.add_display_info(
          feed_source,
          [functools.partial(_feed_display_func, feed_type='FEED_TYPE_BOOK')])
      bid0p_view = CDSView(
          source=feed_source,
          filters=[BooleanFilter(pd.notnull(feed_source.data['bid0p']).tolist())])
      ask0p_view = CDSView(
          source=feed_source,
          filters=[BooleanFilter(pd.notnull(feed_source.data['ask0p']).tolist())])
      bid0p_step = price_fig.step(
          x='time', y='bid0p', source=feed_source, view=bid0p_view, mode='after',
          line_color='yellow', muted_alpha=muted_alpha)
      ask0p_step = price_fig.step(
          x='time', y='ask0p', source=feed_source, view=ask0p_view, mode='after',
          line_color='purple', muted_alpha=muted_alpha)
      feed_legend_items = [('bid0', [bid0p_step]), ('ask0', [ask0p_step])]
    elif feed_type == 'mid':
      price_displayer.add_display_info(
          feed_source,
          [functools.partial(_feed_display_func, feed_type='FEED_TYPE_BOOK')])
      midp_view = CDSView(
          source=feed_source,
          filters=[BooleanFilter(pd.notnull(feed_source.data['midp']).tolist())])
      mid0p = price_fig.step(
          x='time', y='midp', source=feed_source, view=midp_view, mode='after',
          muted_alpha=muted_alpha)
      feed_legend_items = [('mid', [mid0p])]
    elif feed_type == 'trade':
      price_displayer.add_display_info(
          feed_source,
          [functools.partial(_feed_display_func, feed_type='FEED_TYPE_TRADE')])
      tradep_view = CDSView(
          source=feed_source,
          filters=[BooleanFilter(pd.notnull(feed_source.data['tradep']).tolist())])
      tradep = price_fig.step(
          x='time', y='tradep', source=feed_source, view=tradep_view, mode='after',
          muted_alpha=muted_alpha)
      feed_legend_items = [('trade', [tradep])]
    elif feed_type is None:
      pass
    else:
      raise ValueError('Unknown feed type: %s' % feed_type)
  legend = Legend(
      items=order_legend_items+feed_legend_items,
      click_policy='mute', orientation='horizontal', location='top_left')
  price_fig.add_layout(legend, 'above')

  buy_order_price_view = CDSView(
      source=selected_source, filters=[buy_filter, submit_filter])
  sell_order_price_view = CDSView(
      source=selected_source, filters=[sell_filter,  submit_filter])
  buy_fill_price_view = CDSView(
      source=selected_source, filters=[buy_filter, fill_filter])
  sell_fill_price_view = CDSView(
      source=selected_source, filters=[sell_filter, fill_filter])
  buy_cancel_view = CDSView(
      source=selected_source, filters=[buy_filter, cancel_filter])
  sell_cancel_view = CDSView(
      source=selected_source, filters=[sell_filter, cancel_filter])
  price_fig.circle(
      x='event_time', y='order_price', source=selected_source, view=buy_order_price_view,
      line_color='green', line_width=2, fill_color='white')
  price_fig.circle(
      x='event_time', y='order_price', source=selected_source,
      view=sell_order_price_view, line_color='red', line_width=2, fill_color='white')
  price_fig.circle(
      x='event_time', y='fill_price', source=selected_source, view=buy_fill_price_view,
      line_color='green', line_width=2, fill_color='green')
  price_fig.circle(
      x='event_time', y='fill_price', source=selected_source, view=sell_fill_price_view,
      line_color='red', line_width=2, fill_color='red')
  price_fig.segment(
      x0='pending_time', y0='order_price', x1='event_time', y1='fill_price',
      source=selected_source, view=buy_fill_price_view,
      line_color='black', line_width=2)
  price_fig.segment(
      x0='pending_time', y0='order_price', x1='event_time', y1='fill_price',
      source=selected_source, view=sell_fill_price_view,
      line_color='black', line_width=2)
  price_fig.x(
      x='event_time', y='order_price', source=selected_source, view=buy_cancel_view,
      line_color='green', line_width=2)
  price_fig.x(
      x='event_time', y='order_price', source=selected_source, view=sell_cancel_view,
      line_color='red', line_width=2)
  price_fig.segment(
      x0='pending_time', y0='order_price', x1='event_time', y1='order_price',
      source=selected_source, view=buy_cancel_view, line_color='black', line_width=2)
  price_fig.segment(
      x0='pending_time', y0='order_price', x1='event_time', y1='order_price',
      source=selected_source, view=sell_cancel_view, line_color='black', line_width=2)
  return CompositeFigure(price_displayer, price_fig)


def _gen_indicator_displayer(*, indicator, source):
  indicator_displayer = DataSourceDisplayer()
  if indicator in (
      'order_pending_duration', 'order_life_duration', 'num_working_orders'):
    indicator_displayer.add_display_info(
        source,
        [functools.partial(_order_display_func, columns=['event_timestamp', indicator])])
  elif indicator == 'position':
    funcs = [
        functools.partial(
            _order_display_func,
            columns=['event_timestamp', 'position'],
            search_valid_column='position'),
        functools.partial(
            _order_display_func,
            columns=['event_timestamp', 'position_from_fill'],
            search_valid_column='position_from_fill')
    ]
    indicator_displayer.add_display_info(source, funcs)
  elif indicator == 'pnl':
    funcs = [
        functools.partial(
            _order_display_func,
            columns=['event_timestamp', 'pnl_from_fill'],
            search_valid_column='pnl_from_fill')
    ]
    indicator_displayer.add_display_info(source, funcs)
  else:
    logging.error('Unknown indicator: %s' % indicator)
  return indicator_displayer


def _gen_indicator_diagram(
    *, indicator, source, tools, plot_height, plot_width, x_range):
  figure = bokeh.plotting.figure(
      plot_height=plot_height, plot_width=plot_width,
      x_range=x_range, x_axis_type='datetime', tools=tools,
      output_backend='webgl')
  if indicator == 'order_pending_duration':
    order_pending_duration_view = CDSView(
        source=source,
        filters=[BooleanFilter(
            pd.notnull(source.data['order_pending_duration']).tolist())])
    figure.vbar(
        x='event_time', top='order_pending_duration', bottom=0, source=source,
        view=order_pending_duration_view)
    figure.circle(
        x='event_time', y='order_pending_duration',
        fill_color="grey", fill_alpha=0.05, line_color=None,
        source=source, view=order_pending_duration_view)
    figure.yaxis.axis_label = 'order_pending_time'
  elif indicator == 'order_life_duration':
    order_life_duration_view = CDSView(
        source=source,
        filters=[BooleanFilter(
            pd.notnull(source.data['order_life_duration']).tolist())])
    figure.vbar(
        x='event_time', top='order_life_duration', bottom=0, source=source,
        view=order_life_duration_view)
    figure.circle(
        x='event_time', y='order_life_duration',
        fill_color="grey", fill_alpha=0.05, line_color=None,
        source=source, view=order_life_duration_view)
    figure.yaxis.axis_label = 'order_life_time'
  elif indicator == 'num_working_orders':
    num_working_orders_view = CDSView(
        source=source,
        filters=[BooleanFilter(
            pd.notnull(source.data['num_working_orders']).tolist())])
    figure.step(
        x='event_time', y='num_working_orders', source=source,
        view=num_working_orders_view, mode='after')
    figure.circle(
        x='event_time', y='num_working_orders',
        fill_color="grey", fill_alpha=0.05, line_color=None,
        source=source, view=num_working_orders_view)
    figure.yaxis.axis_label = 'working_order'
  elif indicator == 'position':
    position_view = CDSView(
        source=source,
        filters=[BooleanFilter(pd.notnull(source.data['position']).tolist())])
    position_from_fill_view = CDSView(
        source=source,
        filters=[BooleanFilter(
            pd.notnull(source.data['position_from_fill']).tolist())])
    figure.step(
        x='event_time', y='position', source=source, view=position_view, mode='after')
    figure.step(
        x='event_time', y='position_from_fill',
        source=source, view=position_from_fill_view,
        mode='after', line_color='red', line_dash='dashed')
    figure.yaxis.axis_label = 'position'
  elif indicator == 'pnl':
    pnl_from_fill_view = CDSView(
        source=source,
        filters=[BooleanFilter(pd.notnull(source.data['pnl_from_fill']).tolist())])
    # the reason to use step instead of line is because line cannot work well with view
    # https://github.com/bokeh/bokeh/issues/7070
    figure.step(
        x='event_time', y='pnl_from_fill',
        source=source, view=pnl_from_fill_view, mode='after')
    figure.yaxis.axis_label = 'pnl'
  else:
    raise ValueError('Unknown indicator: %s' % indicator)
  return figure


def _gen_range_diagram(*, source, tools, plot_height, plot_width, x_range):
  select_fig = bokeh.plotting.figure(
      plot_height=int(plot_height), plot_width=int(plot_width),
      x_axis_type="datetime", tools="", toolbar_location=None,
      background_fill_color="#efefef")
  range_tool = RangeTool(x_range=x_range)
  range_tool.overlay.fill_color = "navy"
  range_tool.overlay.fill_alpha = 0.2
  #select_fig.circle(
  #    x='event_time', y='order_price', source=source, view=buy_order_price_view)
  select_fig.ygrid.grid_line_color = None
  select_fig.add_tools(range_tool)
  select_fig.toolbar.active_multi = range_tool


def _gen_order_plot(product, strat_log_data, feed_data, feed_type, start_time, end_time):
  strat_log_df = strat_log_data
  source = ColumnDataSource(strat_log_df, name='strategy_log')
  selected_source = ColumnDataSource(pd.DataFrame(columns=strat_log_df.columns))
  feed_sources = {
      product: ColumnDataSource(feed_df, name='feed_%s' % product.full_symbol)
      for product, feed_df in feed_data.items()
  }
  goto_widget = TextInput(
    title='goto', orientation='horizontal',
    value=start_time.strftime('%Y-%m-%d %H:%M:%S'))
  goto_button = Button(label='Goto', button_type='success')
  prev_button = Button(label='Prev', button_type='success')
  next_button = Button(label='Next', button_type='success')
  hybrid_iterator = HybridDataSourceIterator(
    index_threshold=to_timestamp_int(start_time))
  hybrid_iterator.add_data_source(source, 'event_timestamp')
  if feed_type is not None:
    for feed_source in feed_sources.values():
      hybrid_iterator.add_data_source(feed_source, 'timestamp')
  indicator_widget = MultiChoice(
      title='indicator',
      options=['order_pending_duration', 'order_life_duration',
               'num_working_orders', 'position', 'pnl'])
  feed_widget = Select(title='feed', options=['', 'bbo', 'mid', 'trade'])
  tools = "pan,box_select,xwheel_zoom,reset"
  plot_height = 600
  plot_width = 1200
  indicator_figure_universe = {}
  price_composites = []
  indicator_composites = []

  price_composite = _gen_order_sequence_diagram(
      source=source, selected_source=selected_source,
      feed_sources=feed_sources, feed_type=feed_type,
      tools=tools, plot_height=plot_height, plot_width=plot_width,
      start_time=start_time, end_time=end_time)
  price_composites.append(price_composite)
  figure_layout = bokeh.layouts.column(
      children=[price_composite.displayer.widget, price_composite.figure])

  date_format = '%m-%d %H:%M:%S.%3N'
  col_names = [
      'event_time', 'internal_order_id', 'external_order_id', 'type', 'order_price',
      'order_qty', 'order_side', 'order_state', 'fill_price', 'fill_qty',
      'num_working_orders']
  columns = []
  for col_name in col_names:
    if col_name == 'event_time':
      column = TableColumn(
          field=col_name, title=col_name,formatter=DateFormatter(format=date_format),
          width=10)
    else:
      column = TableColumn(field=col_name, title=col_name, width=10)
    columns.append(column)
  table_source = ColumnDataSource(pd.DataFrame(columns=col_names))
  data_table = DataTable(
      source=table_source, columns=columns, css_classes=['my_table'],
      sortable=True, min_width=int(1.3*plot_width), width_policy='min',
      fit_columns=True)
  wo_table_source = ColumnDataSource(pd.DataFrame(columns=col_names))
  wo_table = DataTable(
      source=wo_table_source, columns=columns, css_classes=['my_table'],
      sortable=True, min_width=data_table.min_width, width_policy='min',
      fit_columns=True)
  style = Div(text="""
<style>
.my_table{
  font-family: 'Roboto Mono', Courier, monospace !important;
}
</style>
""")

  def on_indicator_change(attribute, old, new):
    new_indicators = [indicator for indicator in set(new) if indicator not in
                      indicator_figure_universe]
    indicator_displayers = {}
    new_indicator_figures = {}
    for indicator in new:
      try:
        indicator_displayer = _gen_indicator_displayer(
            indicator=indicator, source=source)
        indicator_displayers[indicator] = indicator_displayer
      except Exception:
        logging.error('Fail to gen indicator displayer: %s' % indicator)
        logging.error(traceback.format_exc())
    for indicator in new_indicators:
      try:
        indicator_figure = _gen_indicator_diagram(
            indicator=indicator, source=source,
            tools=tools, plot_height=int(plot_height*0.25), plot_width=plot_width,
            x_range=price_composite.figure.x_range)
        new_indicator_figures[indicator] = indicator_figure
      except Exception:
        logging.error('Fail to gen indicator plot: %s' % indicator)
        logging.error(traceback.format_exc())
    indicator_figure_universe.update(
        {indicator: figure for indicator, figure in new_indicator_figures.items()})
    indicator_composites.clear()
    for indicator in new:
      if indicator in indicator_figure_universe:
        indicator_composite = CompositeFigure(
            displayer=indicator_displayers[indicator],
            figure=indicator_figure_universe[indicator])
        indicator_composites.append(indicator_composite)
    refill_figure_layout()

  def refill_figure_layout():
    children = []
    for composite in price_composites + indicator_composites:
      children += composite.unpack()
    figure_layout.children = children

  def on_goto_button_click():
    dt = _str_to_datetime(goto_widget.value)
    hybrid_iterator.update_by_index_threshold(to_timestamp_int(dt))

  def on_prev_button_click():
    hybrid_iterator.go_next(reverse=True)

  def on_next_button_click():
    hybrid_iterator.go_next()

  def on_feed_change(attribute, old, new):
    old = _parse_feed_type(old)
    new = _parse_feed_type(new)
    if old is None and new is not None:
      for feed_source in feed_sources.values():
        hybrid_iterator.add_data_source(feed_source, 'timestamp')
    if old is not None and new is None:
      for feed_source in feed_sources.values():
        hybrid_iterator.remove_data_source(feed_source)
    price_composite = _gen_order_sequence_diagram(
        source=source, selected_source=selected_source,
        feed_sources=feed_sources, feed_type=new,
        tools=tools, plot_height=plot_height, plot_width=plot_width,
        start_time=start_time, end_time=end_time)
    price_composites.clear()
    price_composites.append(price_composite)
    refill_figure_layout()

  def on_source_change(attribute, old, new):
    if len(new) > 0:
      internal_order_ids = set([source.data['internal_order_id'][idx] for idx in new])
      table_source.data = strat_log_df[strat_log_df['internal_order_id'].isin(internal_order_ids)]
    table_source.selected.indices = []

  def on_table_source_change(attribute, old, new):
    if len(new) > 0:
      internal_order_ids = set()
      for idx in new:
        working_orders = table_source.data['working_orders'][idx]
        if pd.notnull(working_orders):
          internal_order_ids.update(working_orders.split(','))
      wo_table_source.data = strat_log_df[strat_log_df['internal_order_id'].isin(internal_order_ids)]
    wo_table_source.selected.indices = []

  def on_wo_table_source_change(attribute, old, new):
    internal_order_ids = set(
        [wo_table_source.data['internal_order_id'][idx] for idx in new])
    data = wo_table_source.to_df()
    selected_source.data = data[
        data['internal_order_id'].isin(internal_order_ids)]

  def on_iterator_threshold_change(attribute, old, new):
    span_name = 'span'
    location = datetime.datetime.utcfromtimestamp(new / 10**9)
    goto_widget.value = location.strftime('%Y-%m-%d %H:%M:%S.%f')
    for child in figure_layout.children:
      if not isinstance(child, Figure):
        continue
      renderers = child.select(name=span_name)
      if len(renderers) == 0:
        current_time_span = Span(
            name=span_name, location=location, dimension='height',
            line_color='black', line_dash='dashed', line_width=3)
        child.add_layout(current_time_span)
      else:
        assert len(renderers) == 1
        renderers[0].location = location

  def on_iterator_selected_change(attribute, old, new):
    for composite in price_composites + indicator_composites:
      composite.displayer.on_iterator_updated(new)

  indicator_widget.on_change('value', on_indicator_change)
  goto_button.on_click(on_goto_button_click)
  prev_button.on_click(on_prev_button_click)
  next_button.on_click(on_next_button_click)
  feed_widget.on_change('value', on_feed_change)
  source.selected.on_change('indices', on_source_change)
  table_source.selected.on_change('indices', on_table_source_change)
  wo_table_source.selected.on_change('indices', on_wo_table_source_change)
  hybrid_iterator.on_change('threshold', on_iterator_threshold_change)
  hybrid_iterator.on_change('source', on_iterator_selected_change)
  table_layout = bokeh.layouts.column(data_table, wo_table, style)

  layout = bokeh.layouts.column(
      bokeh.layouts.row(indicator_widget, goto_widget, goto_button,
                        prev_button, next_button, feed_widget),
      bokeh.layouts.row(figure_layout, table_layout))
  return layout
  # bokeh.io.show(layout)


def _load_order_event(
    *, root_dir, start_time, end_time, strategy_name, machine, account_request, symbol): 
  data = pd.DataFrame()
  for td in pd.date_range(start_time, end_time):
    try:
      data_dir = get_order_event_data_dir(
          root_dir=root_dir,
          start_date=td,
          end_date=td,
          strategy_name=strategy_name,
          machine=machine,
          account_request=account_request,
          symbol=symbol)
      files = [filename for filename in os.listdir(data_dir) if
               filename.endswith('.json')]
      assert len(files) == 1
      with open(os.path.join(data_dir, files[0])) as fd:
        data = data.append(pd.DataFrame(json.load(fd)))
    except Exception:
      logging.error('Fail to load order event: %s' % td.date())
      logging.error(traceback.format_exc())
  data = data[(data['event_timestamp'] >= to_timestamp_int(start_time))
              & (data['event_timestamp'] <= to_timestamp_int(end_time))]
  to_datetime_func = \
      lambda x: None if pd.isnull(x) else datetime.datetime.utcfromtimestamp(x / 10**9)
  data['event_time'] = data['event_timestamp'].apply(to_datetime_func)
  data['pending_time'] = data['pending_timestamp'].apply(to_datetime_func)
  return data


def _load_feed(*, feed_dir, start_time, end_time, machine, products):
  replayer = FeedReplayer(products, start_time, end_time)
  abase.run_from_feed_cache_base(
      products,
      replayer.on_book_reset,
      start_time.replace(tzinfo=None),
      end_time.replace(tzinfo=None),
      machine,
      ['1', '2'],
      feed_cache_dir=feed_dir)
  feed = {product : pd.DataFrame(feed) for product, feed in replayer.feed.items()
          if len(feed) > 0}
  return feed


class CompositeFigure(object):
  def __init__(self, displayer, figure):
    assert isinstance(displayer, DataSourceDisplayer)
    assert isinstance(figure, Figure)
    self.displayer = displayer
    self.figure = figure

  def unpack(self):
    return [self.displayer.widget, self.figure]


class FeedReplayer(object):
  def __init__(self, products, start_time=None, end_time=None):
    self._products = products
    self._feed = {product: [] for product in products}

  @property
  def feed(self):
    return self._feed

  def on_book_reset(self, book_builder_name, book_builder):
    for product in self._products:
      book_builder.subscribe(product, functools.partial(self.on_book, product))
      if hasattr(book_builder, 'subscribe_trade'):
        book_builder.subscribe_trade(product, functools.partial(self.on_trade, product))

  def on_book(self, product, book):
    bid0p = None
    bid0q = None
    ask0p = None
    ask0q = None
    if book.has_bid():
      bid0p = book.bid0().price
      bid0q = book.bid0().qty
    if book.has_ask():
      ask0p = book.ask0().price
      ask0q = book.ask0().qty
    midp = (bid0p + ask0p) / 2 if bid0p is not None and ask0p is not None else None
    timestamp = book.timestamp
    self._feed[product].append({
        'timestamp': timestamp,
        'time': datetime.datetime.utcfromtimestamp(timestamp / 10**9),
        'feed_type': 'FEED_TYPE_BOOK',
        'bid0p': bid0p,
        'bid0q': bid0q,
        'ask0p': ask0p,
        'ask0q': ask0q,
        'midp': midp
      })

  def on_trade(self, product, trade):
    timestamp = trade.timestamp
    self._feed[product].append({
        'timestamp': trade.timestamp,
        'time': datetime.datetime.utcfromtimestamp(timestamp / 10 ** 9),
        'feed_type': 'FEED_TYPE_TRADE',
        'trade_side': feedpb.TradeSide.Name(trade.side),
        'tradep': trade.price,
        'tradeq': trade.qty,
      })


class DataSourceDisplayInfo(object):
  def __init__(self, display_funcs):
    assert isinstance(display_funcs, list)
    self._display_funcs = display_funcs
    self._text = None

  @property
  def display_funcs(self):
    return self._display_funcs

  @property
  def text(self):
    return self._text

  def set_text(self, text):
    self._text = text


class DataSourceDisplayer(object):
  def __init__(self):
    self._widget = Div()
    self._display_infos = {}

  @property
  def widget(self):
    return self._widget

  def _update_display_info(self):
    text = ''
    for name, info in self._display_infos.items():
      text += '' if info.text is None else '<br>%s' % info.text
    self._widget.text = text

  def on_iterator_updated(self, iterator):
    assert isinstance(iterator, DataSourceIterator)
    for name, info in self._display_infos.items():
      if name == iterator.name:
        text = []
        for func in info.display_funcs:
          text.append(func(iterator))
        info.set_text(', '.join(text))
    self._update_display_info()

  def add_display_info(self, data_source, display_funcs):
    name = data_source.name
    assert name is not None and name not in self._display_infos
    self._display_infos[name] = DataSourceDisplayInfo(display_funcs)
    self._update_display_info()

  def remove_data_source(self, data_source):
    name = data_source.name
    assert name is not None and name not in self._display_infos
    del self._display_infos[name]
    self._update_display_info()


class DataSourceIterator(object):
  def __init__(self, data_source, index_column, index_threshold):
    assert isinstance(data_source, ColumnDataSource)
    assert index_column in data_source.data.keys()
    self._name = data_source.name
    self._data = data_source.to_df()
    self._index_column = index_column
    self._index_threshold = None
    self.update_by_index_threshold(index_threshold)

  @property
  def data(self):
    return self._data

  @property
  def name(self):
    return self._name

  @property
  def row_idx(self):
    return self._row_idx

  @property
  def index_column(self):
    return self._index_column

  def set_row_idx(self, row_idx):
    self._row_idx = row_idx

  def update_by_index_threshold(self, index_threshold):
    if index_threshold != self._index_threshold:
      filtered = self._data[self._data[self._index_column] <= index_threshold]
      if len(filtered) > 0:
        row_idx = filtered.index[-1]
      else:
        row_idx = None
      self._index_threshold = index_threshold
      self._row_idx = row_idx

  def get_next_row(self, offset):
    if offset > 0:
      next_row_idx = 0 if self._row_idx is None else self._row_idx + offset
    else:
      next_row_idx = None if self._row_idx is None else self._row_idx + offset
    if next_row_idx is not None:
      if next_row_idx < 0 or next_row_idx >= len(self._data):
        next_row_idx = None
    next_row = None if next_row_idx is None else self._data.iloc[next_row_idx]
    return next_row_idx, next_row


class HybridDataSourceIterator(object):
  def __init__(self, index_threshold):
    assert index_threshold is not None
    self._index_threshold = index_threshold
    self._selected = _DataSourceInfo(name=None, row_idx=None)
    self._data_source_infos = {}
    self._changed_data_sources = []
    self._callbacks = {}

  def _get_next_row_offset(self, data_source, reverse=False):
    if reverse:
      if self._selected.name is not None and self._selected.name == data_source.name:
        return -1
      else:
        return 0
    else:
      return 1

  def _update_selected(self, selected=None):
    if selected is None:
      index_value_dict = {}
      for name, data_source in self._data_source_infos.items():
        curr_row_idx, curr_row = data_source.get_next_row(0)
        if curr_row is not None:
          index_value_dict[name] = {
              'row_idx': curr_row,
              'index_value': curr_row[data_source.index_column]
          }
      if len(index_value_dict) > 0:
        source_name = max(
            index_value_dict.keys(), key=(lambda k: index_value_dict[k]['index_value']))
        self._set_selected(_DataSourceInfo(
            name=source_name,
            row_idx=self._data_source_infos[source_name].row_idx))
    else:
      self._set_selected(selected)

  def _set_index_threshold(self, threshold):
    prev_threshold = self._index_threshold
    if prev_threshold != threshold:
      self._index_threshold = threshold
      self._run_callback('threshold', prev_threshold, threshold)

  def _set_selected(self, selected):
    assert isinstance(selected, _DataSourceInfo)
    prev_selected = self._selected
    if prev_selected != selected:
      self._selected = selected
      self._changed_data_sources.append(self._data_source_infos[self._selected.name])

  def _run_callback(self, attribute, old, new):
    func = self._callbacks.get(attribute)
    if func is not None:
      func(attribute, old, new)

  def on_change(self, attribute, callback):
    assert attribute in ('source', 'threshold')
    self._callbacks[attribute] = callback

  def _publish_souce_change(self):
    source_set = set()
    for source in self._changed_data_sources:
      if source.name not in source_set:
        source_set.add(source.name)
        self._run_callback('source', None, source)
    self._changed_data_sources.clear()

  def update_by_index_threshold(self, threshold):
    for key, data_source_info in self._data_source_infos.items():
      data_source_info.update_by_index_threshold(threshold)
    self._set_index_threshold(threshold)
    self._update_selected()
    self._publish_souce_change()

  def go_next(self, reverse=False):
    index_value_dict = {}
    for name, data_source in self._data_source_infos.items():
      offset = self._get_next_row_offset(data_source, reverse=reverse)
      next_row_idx, next_row = data_source.get_next_row(offset)
      # must move first if offset < 0
      if offset < 0:
        data_source.set_row_idx(next_row_idx)
        self._changed_data_sources.append(data_source)
      if next_row is not None:
        index_value_dict[name] = {
            'row_idx': next_row_idx,
            'index_value': next_row[data_source.index_column]
        }
    if len(index_value_dict) > 0:
      func = max if reverse else min
      source_name = func(
          index_value_dict.keys(), key=(lambda k: index_value_dict[k]['index_value']))
      selected_data_source = self._data_source_infos[source_name]
      selected_row_idx = index_value_dict[source_name]['row_idx']
      selected_data_source.set_row_idx(selected_row_idx)
      self._set_index_threshold(index_value_dict[source_name]['index_value'])
      self._update_selected(selected=_DataSourceInfo(
          name=source_name,
          row_idx=selected_row_idx))
      self._publish_souce_change()

  def add_data_source(self, data_source, index_column):
    name = data_source.name
    assert name is not None and name not in self._data_source_infos
    self._data_source_infos[name] = DataSourceIterator(
        data_source, index_column, self._index_threshold)
    self._update_selected()
    self._publish_souce_change()

  def remove_data_source(self, data_source):
    name = data_source.name
    assert name is not None and name in self._data_source_infos
    del self._data_source_infos[name]
    self._update_selected()
    self._publish_souce_change()

  def get_row(self, data_source):
    return self._data_source_infos[data_source.name].row_idx


class OrderInfoPicker(object):
  def __init__(self):
    widgets = {}
    self._titles = [
        'start_time', 'end_time', 'strategy_name', 'machine', 'account', 'symbol']
    for title in self._titles:
      if title in ['start_time', 'end_time', 'strategy_name']:
        value = None
        if title in ('start_time', 'end_time'):
          prev_dt = datetime.datetime.utcnow() - datetime.timedelta(days=1)
          value = datetime.datetime(prev_dt.year, prev_dt.month, prev_dt.day).strftime(
              '%Y-%m-%d %H:%M:%S')
        widget = TextInput(
            title=title, orientation='horizontal', name=title, value=value)
        widget.on_change('value', self._update_order_info)
      else:
        widget = Select(title=title, name=title)
      widgets[title] = widget
    self._widgets = widgets

  @property
  def titles(self):
    return self._titles

  @property
  def widgets(self):
    return list(self._widgets.values())

  def get_widget(self, title):
    return self._widgets[title]

  def _update_order_info(self, attr, old, new):
    start_date = None
    end_date = None
    strategy_name = None
    for widget in self.widgets:
      try:
        if widget.title == 'start_time' and widget.value is not None:
          start_date = _str_to_datetime(widget.value).date()
        if widget.title == 'end_time' and widget.value is not None:
          end_date = _str_to_datetime(widget.value).date()
        if widget.title == 'strategy_name':
          strategy_name = widget.value
      except:
        pass
    fetch_cols = [
        'strategy_name', 'machine', 'market_type', 'exchange', 'owner', 'symbol']
    if start_date is None or end_date is None or strategy_name is None:
      return
    pta_df = query_pta(start_date, end_date,
                       business_units=['Coin'], pta_type='SYMBOL')[fetch_cols]
    pta_df = pta_df[pta_df['strategy_name'] == strategy_name].drop_duplicates()
    order_info = {}
    order_info['machine'] = pta_df['machine']
    order_info['account'] = ['%s.%s.%s' % (market_type, exchange, owner)
                             for market_type in pta_df['market_type']
                             for exchange in pta_df['exchange']
                             for owner in pta_df['owner']]
    order_info['symbol'] = pta_df['symbol']
    for key, value in order_info.items():
      options = sorted(list(set(value)))
      self._widgets[key].options = options
      self._widgets[key].value = options[0] if len(options) > 0 else None


def bkapp(doc):
  arguments = doc.session_context.request.arguments
  layouts = bokeh.layouts.column()
  status = Div()
  try:
    order_info_picker = OrderInfoPicker()
    _get_param = lambda x: order_info_picker.get_widget(x).value
    for title in order_info_picker.titles:
      argument = arguments[title][0].decode().strip() if title in arguments else None
      if argument is not None:
        order_info_picker.get_widget(title).value = argument
    submit_button = Button(label='Submit', button_type='success')
    code = """
      var args = {};
      for (var i = 0; i < widgets.length; i++) {
        console.log(widgets[i].value);
        if (widgets[i].value !== null) {
          args[widgets[i].name] = widgets[i].value
        }
      }
      const paramString = new URLSearchParams(args);
      window.open(`http://%s:%s/?${paramString.toString()}`, "_self");
    """ % (FLAGS.address, FLAGS.port)
    click_callback = CustomJS(args=dict(widgets=order_info_picker.widgets), code=code)
    submit_button.js_on_click(click_callback)
    control_widget = bokeh.layouts.row(order_info_picker.widgets + [submit_button])
    layouts = bokeh.layouts.column(control_widget, status)

    market_type, exchange, owner = _get_param('account').split('.')
    acct_req = cr_pb2.AccountRequestProto(
        market_type=market_type, exchange=exchange, owner=owner)
    start_time = _str_to_datetime(_get_param('start_time'))
    end_time = _str_to_datetime(_get_param('end_time'))
    product = generate_product_from_str2(
        market_type, exchange, None, _get_param('symbol'),
        start_time.replace(tzinfo=None))
    fetch_func_list = [
        functools.partial(_load_order_event,
                          root_dir=FLAGS.order_event_dir,
                          start_time=start_time,
                          end_time=end_time,
                          strategy_name=_get_param('strategy_name'),
                          machine=_get_param('machine'),
                          account_request=acct_req,
                          symbol=_get_param('symbol')),
        functools.partial(_load_feed,
                          feed_dir=FLAGS.feed_dir,
                          start_time=start_time,
                          end_time=end_time,
                          machine='feed-best.ap-northeast-1.aws',
                          products=[product])
    ]
    with concurrent.futures.ThreadPoolExecutor() as executor:
      future_list = [executor.submit(func) for func in fetch_func_list]
    strat_log_data, feed_data = [future.result() for future in future_list]
    plot_layouts = _gen_order_plot(
        product, strat_log_data, feed_data=feed_data, feed_type=None,
        start_time=start_time, end_time=end_time)
    layouts = bokeh.layouts.column(layouts, plot_layouts)
  except Exception:
    status.text = '<p>Fail to generate order plot</p>'
    logging.error(traceback.format_exc())
  doc.add_root(layouts)
  """
  strat_log_root_dir= '~/data/strat_proto_log'
  machine = 'strategy-43.ap-northeast-1.huobi'
  start_time = datetime.datetime(2020, 8, 5, tzinfo=pytz.UTC)
  end_time = datetime.datetime(2020, 8, 5, 0, 59, 59, tzinfo=pytz.UTC)
  strategy_name = 'lm_agg2_huobi_futures_btc'
  """


def main(_):
  for path in (FLAGS.order_event_dir, FLAGS.feed_dir):
    assert path is not None and os.path.exists(path), path
  server = Server({'/': bkapp}, io_loop=IOLoop(), port=FLAGS.port)
  server.start()
  server.io_loop.start()


if __name__ == '__main__':
  abase.define_base_flags()
  abase.define_feed_archive_flags()
  flags.DEFINE_string('order_event_dir', None, '')
  flags.DEFINE_string('feed_dir', '/remote/iosg/coin-2/feed_cache', '')
  flags.DEFINE_string('address', None, '')
  flags.DEFINE_integer('port', 5006, '')

  app.run(main)
